# %% Import relevant modules
from matplotlib import pyplot
from scipy.stats import zscore
import pandas as pd
import numpy as np
import warnings
import os
from sklearn.decomposition import PCA

warnings.filterwarnings("ignore")
# %% Loading data
path = "./EquityData"
index = '000905.XSHG'
#{'000300.XSHG', '000852.XSHG', '000905.XSHG', 'not in any index'}
meta_data = pd.read_csv('./A_imputer.csv')
meta_data.drop(columns=['Unnamed: 0', 'close_y'], axis=1, inplace=True) # 剔除序号列和重复close

# 取各市场的数据
meta_data_300 = meta_data[meta_data['index'] == '000300.XSHG'].sort_values(['time', 'code']).reset_index(drop=True)
meta_data_300.drop(columns=['index'], inplace=True, axis=1)
meta_data_300 = meta_data_300.sort_values(['time', 'code']).reset_index(drop=True)

meta_data_500 = meta_data[meta_data['index'] == '000852.XSHG'].sort_values(['time', 'code']).reset_index(drop=True)
meta_data_500.drop(columns=['index'], inplace=True, axis=1)
meta_data_500 = meta_data_500.sort_values(['time', 'code']).reset_index(drop=True)

meta_data_1000 = meta_data[meta_data['index'] == '000905.XSHG'].sort_values(['time', 'code']).reset_index(drop=True)
meta_data_1000.drop(columns=['index'], inplace=True, axis=1)
meta_data_1000 = meta_data_1000.sort_values(['time', 'code']).reset_index(drop=True)

meta_data_all = meta_data.sort_values(['time', 'code']).reset_index(drop=True)
meta_data_all.drop(columns=['index'], inplace=True, axis=1)
meta_data_all = meta_data_all.sort_values(['time', 'code']).reset_index(drop=True)

# 预测目标：下一期收盘价
close = meta_data[['time', 'code', 'close_x']]
del meta_data

# 所有参数后移一天
def return_shift_one(group):
    # group[group.columns[2:]] = group[group.columns[2:]].shift(1)
    group[group.columns[2:]] = group[group.columns[2:]].shift(1)
    return group

def shift_all_columns(meta_data):
    meta_data = meta_data.groupby(['code']).apply(return_shift_one).reset_index(
            drop=True).sort_values(['time', 'code']).reset_index(drop=True)
    meta_data.dropna(axis=0, inplace=True)
    return meta_data

meta_data_300 = shift_all_columns(meta_data_300)
meta_data_500 = shift_all_columns(meta_data_500)
meta_data_1000 = shift_all_columns(meta_data_1000)
meta_data_all = shift_all_columns(meta_data_all)


#% PCA
# 对除time和code外的所有特征PCA
def pca_meta_data(group, n_compents=25):
    pca = PCA(n_components=n_compents)
    group.reset_index(inplace=True, drop=True)
    group_pca = pd.DataFrame(pca.fit_transform(group[group.columns[2:]]))
    group_pca = pd.concat([group[group.columns[:2]], group_pca], axis=1)
    return group_pca

def apply_pca(meta_data, n_compents=25):
    meta_data = meta_data.groupby(['time']).apply(pca_meta_data).reset_index(
    drop=True).sort_values(['time', 'code']).reset_index(drop=True)
    meta_data = pd.merge(meta_data, close, on=['time', 'code'], how='left')
    meta_data.to_csv('./' + str(len(meta_data)) + '.csv')
    return 0

apply_pca(meta_data_300)
apply_pca(meta_data_500)
apply_pca(meta_data_1000)
apply_pca(meta_data_all)